package swordoffer.offer40;

public class QuickSelect {
    public static void main(String[] args) {
        int[] nums = {4, 5, 1, 6, 2, 7, 3, 8};
    }

    public int[] getLeastNumbers(int[] arr, int k) {
        // 特殊情况
        if (arr.length == 0 || k == 0) {
            return new int[0];
        }
        if (k < 0 || k > arr.length) {
            return arr;
        }
        // 通过快排排好第 K 小的数（下标为k - 1），则左边的数就是另外最小的 K - 1 个数
        int index = quickSelect(arr, 0, arr.length - 1, k - 1);
        // 返回的 index 是第 K 小的数的下标
        int[] res = new int[index + 1];
        // 第 K 小及其左边的数就是所要的答案
        for (int i = 0; i <= index; i++) {
            res[i] = arr[i];
        }
        return res;
    }

    /*
        注意这个函数应该返回 **下标**
     */
    public int quickSelect(int[] arr, int lo, int hi, int target) {
        if (lo >= hi) {
            return lo;
        }
        int pivotIndex = partition(arr, lo, hi);
        if (pivotIndex == target) {
            return pivotIndex;
        } else if (pivotIndex > target) {
            return quickSelect(arr, lo, pivotIndex - 1, target);
        } else {
            return quickSelect(arr, pivotIndex + 1, hi, target);
        }
    }

    public int partition(int[] arr, int lo, int hi) {
        int pivot = arr[lo];
        int left = lo;
        int right = hi;
        while (left < right) {
            while (left < right && arr[right] > pivot) {
                right--;
            }
            while (left < right && arr[left] <= pivot) {
                left++;
            }
            if (left < right) {
                swap(arr, left, right);
            }
        }
        arr[lo] = arr[left];
        arr[left] = pivot;
        return left;
    }

    public void swap(int[] nums, int i, int j) {
        int temp = nums[i];
        nums[i] = nums[j];
        nums[j] = temp;
    }
}
